from pinecone import Pinecone, ServerlessSpec
import matplotlib.pyplot as plt
import numpy as np
from collections import Counter
from sklearn.datasets import load_digits

# 设置中文字体，解决中文显示问题
plt.rcParams["font.family"] = ["SimHei"]

# 初始化Pinecone客户端
pc = Pinecone(api_key="pcsk_dqdo1_EP62eXyBSFyv5dQWZiMJZybi6YAkjrT8cYaobXeTkRDvfnbvxY9y8RCdZkajBb1")
index_name = "mnist-index"

# 检查并清除现有索引
existing_indexes = pc.list_indexes()
if any(index['name'] == index_name for index in existing_indexes):
    print(f"删除现有索引 '{index_name}'...")
    pc.delete_index(index_name)


# 创建新索引
print(f"创建索引 '{index_name}'...")
pc.create_index(
    name=index_name,
    dimension=64,
    metric="euclidean",
    spec=ServerlessSpec(cloud="aws", region="us-east-1")
)
# 等待索引 ready
import time
while True:
    desc = pc.describe_index(index_name)
    if desc['status']['ready']:
        break
    print("索引未就绪，等待中...")
    time.sleep(2)
index = pc.Index(index_name)
print(f"已连接到索引 '{index_name}'")


# 加载MNIST数据集
digits = load_digits(n_class=10)
X = digits.data.astype(np.float32)  # 确保为 float32
y = digits.target
print(f"加载了 {len(X)} 个MNIST样本")

# 准备并上传向量
vectors = []
for i in range(len(X)):
    vectors.append((str(i), X[i].tolist(), {"label": int(y[i])}))



# 分批上传（必须用 vectors= 关键字参数）
batch_size = 1000
for i in range(0, len(vectors), batch_size):
    batch = vectors[i:i + batch_size]
    upsert_response = index.upsert(vectors=batch)
    print("upsert_response:", upsert_response)
    try:
        count = upsert_response.upserted_count
    except Exception:
        count = upsert_response if isinstance(upsert_response, int) else str(upsert_response)
    print(f"已上传第 {i//batch_size + 1} 批向量，状态: {count} 个成功")
    # 调试：每次 upsert 后打印索引描述和索引列表
    print("describe_index:", pc.describe_index(index_name))
    print("list_indexes:", pc.list_indexes())


# 检查索引中的向量数量

# 检查索引中的向量数量，并打印调试信息
stats = index.describe_index_stats()
print("stats:", stats)
try:
    total_count = stats.total_vector_count
except Exception:
    total_count = stats['total_vector_count'] if isinstance(stats, dict) and 'total_vector_count' in stats else str(stats)
print(f"索引中的向量总数: {total_count}")

    
# 如果索引中没有向量，提示错误并退出
if int(total_count) == 0:
    print("错误：向量上传失败，请检查API密钥和网络连接")
else:
    # 创建查询图像（数字3）
    digit_3 = np.array([
        [0, 0, 255, 255, 255, 255, 0, 0],
        [0, 0, 0, 0, 0, 255, 0, 0],
        [0, 0, 0, 0, 0, 255, 0, 0],
        [0, 0, 0, 255, 255, 255, 0, 0],
        [0, 0, 0, 0, 0, 255, 0, 0],
        [0, 0, 0, 0, 0, 255, 0, 0],
        [0, 0, 0, 0, 0, 255, 0, 0],
        [0, 0, 255, 255, 255, 255, 0, 0]
    ])


    # 预处理查询向量（缩放到0-16，并转为float32，与索引数据一致）
    digit_3_normalized = (digit_3 / 16).astype(np.float32)
    query_data = digit_3_normalized.ravel().tolist()

    # 执行查询
    print("正在执行向量查询...")
    results = index.query(
        vector=query_data,
        top_k=11,
        include_metadata=True
    )

    # 处理查询结果
    if not results['matches']:
        print("警告：未找到匹配的向量，可能是查询向量与训练数据差异过大")
        final_prediction = "未知"
    else:
        labels = [match['metadata']['label'] for match in results['matches']]
        for match, label in zip(results['matches'], labels):
            print(f"id: {match['id']}, 距离: {match['score']:.4f}, 标签: {label}")
        
        final_prediction = Counter(labels).most_common(1)[0][0]

    # 显示查询图像和结果
    plt.imshow(digit_3, cmap='gray')
    plt.title(f"预测数字: {final_prediction}", size=15)
    plt.axis('off')
    plt.show()
